<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>KV Cache Size Calculator</title>
    <style>
        * {
            box-sizing: border-box;
        }
        
        body {
            font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, "Helvetica Neue", Arial, sans-serif;
            margin: 0;
            padding: 0;
            background: transparent;
            overflow: hidden;
        }
        
        .calculator-container {
            max-width: 100%;
            margin: 0;
            background: transparent;
            padding: 0;
        }
        
        h1 {
            display: none;
        }
        
        .form-group {
            margin-bottom: 20px;
        }
        
        label {
            display: block;
            font-weight: 600;
            margin-bottom: 8px;
            color: #333;
            font-size: 14px;
        }
        
        select, input {
            width: 100%;
            padding: 12px;
            font-size: 15px;
            border: 2px solid #e0e0e0;
            border-radius: 8px;
            background: white;
            transition: border-color 0.2s;
        }
        
        select:focus, input:focus {
            outline: none;
            border-color: #3898ec;
        }
        
        button {
            width: 100%;
            padding: 14px;
            font-size: 16px;
            font-weight: 600;
            background: #3898ec;
            color: white;
            border: none;
            border-radius: 8px;
            cursor: pointer;
            transition: background 0.2s;
        }
        
        button:hover {
            background: #2578cc;
        }
        
        button:active {
            transform: scale(0.98);
        }
        
        .result {
            margin-top: 24px;
            padding: 20px;
            background: #f0f9ff;
            border-left: 4px solid #3898ec;
            border-radius: 8px;
            display: none;
        }
        
        .result.show {
            display: block;
        }
        
        .result-value {
            font-size: 32px;
            font-weight: 700;
            color: #1a1a1a;
            margin: 8px 0;
        }
        
        .result-label {
            font-size: 14px;
            color: #666;
            margin-bottom: 4px;
        }
        
        .details {
            margin-top: 16px;
            padding-top: 16px;
            border-top: 1px solid #d0e8f7;
            font-size: 13px;
            color: #555;
            line-height: 1.6;
        }
        
        .details strong {
            color: #333;
        }
        
        .error {
            background: #fee;
            border-left-color: #f44;
        }
        
        .loading {
            text-align: center;
            color: #666;
            padding: 20px;
            font-size: 14px;
        }
        
    </style>
</head>
<body>
    <div class="calculator-container">
        <h1>KV Cache Size Calculator</h1>
        
        <div class="loading" id="loading">Loading model configurations...</div>
        
        <div id="calculator-form" style="display: none;">
            <div class="form-group">
                <label for="model">Model:</label>
                <select id="model">
                    <!-- Options will be populated dynamically -->
                </select>
            </div>
            
            <div class="form-group">
                <label for="dtype">Data Type:</label>
                <select id="dtype">
                    <option value="float16">float16 (FP16)</option>
                    <option value="bfloat16" selected>bfloat16 (BF16)</option>
                    <option value="float32">float32 (FP32)</option>
                    <option value="int8">int8 (INT8)</option>
                </select>
            </div>
            
            <div class="form-group">
                <label for="tokens">Number of Tokens:</label>
                <input type="number" id="tokens" placeholder="e.g., 1000" min="1" value="1000">
            </div>
            
            <button onclick="calculateKVCache()">Calculate KV Cache Size</button>
            
            <div id="result" class="result"></div>
            
            <hr style="margin: 40px 0; border: none; border-top: 1px solid #e0e0e0;">
            
            <h2 style="font-size: 18px; margin-bottom: 20px; color: #1a1a1a;">Reverse Calculator: Find Maximum Tokens</h2>
            
            <div class="form-group">
                <label for="model-reverse">Model:</label>
                <select id="model-reverse">
                    <!-- Options will be populated dynamically -->
                </select>
            </div>
            
            <div class="form-group">
                <label for="dtype-reverse">Data Type:</label>
                <select id="dtype-reverse">
                    <option value="float16">float16 (FP16)</option>
                    <option value="bfloat16" selected>bfloat16 (BF16)</option>
                    <option value="float32">float32 (FP32)</option>
                    <option value="int8">int8 (INT8)</option>
                </select>
            </div>
            
            <div class="form-group">
                <label for="gpu-ram">GPU RAM Size (GB):</label>
                <input type="number" id="gpu-ram" placeholder="e.g., 20" min="0.1" step="0.1" value="20">
            </div>
            
            <button onclick="calculateMaxTokens()">Calculate Maximum Tokens</button>
            
            <div id="result-reverse" class="result"></div>
        </div>
    </div>
    
    <script>
        let modelConfigs = {};
        
        // Load model configurations
        async function loadModelConfigs() {
            const urls = [
                'modelconfig.json',
                '_static/modelconfig.json',
                '../_static/modelconfig.json',
                'https://raw.githubusercontent.com/LMCache/LMCache/refs/heads/dev/examples/kv_cache_calculator/modelconfig.json'
            ];
            
            for (const url of urls) {
                try {
                    const response = await fetch(url);
                    if (response.ok) {
                        modelConfigs = await response.json();
                        populateModelDropdown();
                        document.getElementById('loading').style.display = 'none';
                        document.getElementById('calculator-form').style.display = 'block';
                        return;
                    }
                } catch (error) {
                    console.log('Failed to load from:', url);
                }
            }
            
            document.getElementById('loading').innerHTML = 
                '<div class="error">Failed to load model configurations. Please check your connection.</div>';
        }
        
        // Populate model dropdown
        function populateModelDropdown() {
            const select = document.getElementById('model');
            const selectReverse = document.getElementById('model-reverse');
            select.innerHTML = '';
            selectReverse.innerHTML = '';
            
            const collator = new Intl.Collator(undefined, { numeric: true, sensitivity: 'base' });
            const sortedModels = Object.keys(modelConfigs).sort(collator.compare);
            const defaultModel = 'Qwen/Qwen3-8B';
            
            for (const modelName of sortedModels) {
                const option = document.createElement('option');
                option.value = modelName;
                option.textContent = modelName;
                if (modelName === defaultModel) {
                    option.selected = true;
                }
                select.appendChild(option);
                
                const optionReverse = document.createElement('option');
                optionReverse.value = modelName;
                optionReverse.textContent = modelName;
                if (modelName === defaultModel) {
                    optionReverse.selected = true;
                }
                selectReverse.appendChild(optionReverse);
            }
        }
        
        // Calculate KV cache size
        function calculateKVCache() {
            const model = document.getElementById('model').value;
            const tokens = parseInt(document.getElementById('tokens').value);
            const dtype = document.getElementById('dtype').value;
            const resultDiv = document.getElementById('result');
            
            // Validation
            if (isNaN(tokens) || tokens <= 0) {
                resultDiv.className = 'result show error';
                resultDiv.innerHTML = '<div class="result-label">Error</div><div>Please enter a valid number of tokens (greater than 0).</div>';
                return;
            }
            
            const config = modelConfigs[model];
            if (!config) {
                resultDiv.className = 'result show error';
                resultDiv.innerHTML = '<div class="result-label">Error</div><div>Model configuration not found.</div>';
                return;
            }
            
            // Get dtype size
            const dtypeSizes = {
                'float32': 4,
                'float16': 2,
                'bfloat16': 2,
                'int8': 1
            };
            const dtypeSize = dtypeSizes[dtype];
            
            // Calculate based on model architecture
            const isDeepSeek = model.includes('DeepSeek');
            const isQwen3 = model.toLowerCase().includes('qwen/qwen3-');
            
            let totalElements;
            let totalBytes;
            let sizeGB;
            let details = '';
            
            if (isDeepSeek && config.kv_lora_rank && config.qk_rope_head_dim) {
                // DeepSeek V3/R1 architecture
                totalElements = config.num_hidden_layers * tokens * (config.kv_lora_rank + config.qk_rope_head_dim);
                totalBytes = totalElements * dtypeSize;
                sizeGB = totalBytes / (1024 ** 3);
                details = `
                    <strong>Selected Model:</strong> ${model}<br>
                    <strong>Number of Hidden Layers:</strong> ${config.num_hidden_layers}<br>
                    <strong>KV LoRA Rank:</strong> ${config.kv_lora_rank}<br>
                    <strong>QK RoPE Head Dim:</strong> ${config.qk_rope_head_dim}<br>
                    <strong>Data Type Size:</strong> ${dtypeSize} bytes<br>
                    <strong>Total Elements:</strong> ${config.num_hidden_layers} × ${tokens} × (${config.kv_lora_rank} + ${config.qk_rope_head_dim}) = ${totalElements}<br>
                    <strong>Total Bytes:</strong> ${totalElements} × ${dtypeSize} = ${totalBytes} bytes<br>
                    <strong>KV Cache Size:</strong> ${totalBytes} / (1024³) ≈ ${sizeGB.toFixed(4)} GB
                `;
            } else if (isQwen3 && config.head_dim) {
                // Qwen3 with explicit head_dim
                totalElements = 2 * config.num_hidden_layers * tokens * config.num_key_value_heads * config.head_dim;
                totalBytes = totalElements * dtypeSize;
                sizeGB = totalBytes / (1024 ** 3);
                details = `
                    <strong>Selected Model:</strong> ${model}<br>
                    <strong>Hidden Size:</strong> ${config.hidden_size}<br>
                    <strong>Number of Attention Heads:</strong> ${config.num_attention_heads}<br>
                    <strong>Number of Hidden Layers:</strong> ${config.num_hidden_layers}<br>
                    <strong>Number of Key-Value Heads:</strong> ${config.num_key_value_heads}<br>
                    <strong>Head Dimension:</strong> ${config.head_dim}<br>
                    <strong>Data Type Size:</strong> ${dtypeSize} bytes<br>
                    <strong>Total Elements:</strong> 2 × ${config.num_hidden_layers} × ${tokens} × ${config.num_key_value_heads} × ${config.head_dim} = ${totalElements}<br>
                    <strong>Total Bytes:</strong> ${totalElements} × ${dtypeSize} = ${totalBytes} bytes<br>
                    <strong>KV Cache Size:</strong> ${totalBytes} / (1024³) ≈ ${sizeGB.toFixed(4)} GB
                `;
            } else {
                // Standard transformer architecture
                const headSize = config.hidden_size / config.num_attention_heads;
                totalElements = 2 * config.num_hidden_layers * tokens * config.num_key_value_heads * headSize;
                totalBytes = totalElements * dtypeSize;
                sizeGB = totalBytes / (1024 ** 3);
                details = `
                    <strong>Selected Model:</strong> ${model}<br>
                    <strong>Hidden Size:</strong> ${config.hidden_size}<br>
                    <strong>Number of Attention Heads:</strong> ${config.num_attention_heads}<br>
                    <strong>Number of Hidden Layers:</strong> ${config.num_hidden_layers}<br>
                    <strong>Number of Key-Value Heads:</strong> ${config.num_key_value_heads}<br>
                    <strong>Head Size:</strong> ${headSize} (Hidden Size / Attention Heads)<br>
                    <strong>Data Type Size:</strong> ${dtypeSize} bytes<br>
                    <strong>Total Elements:</strong> 2 × ${config.num_hidden_layers} × ${tokens} × ${config.num_key_value_heads} × ${headSize} = ${totalElements}<br>
                    <strong>Total Bytes:</strong> ${totalElements} × ${dtypeSize} = ${totalBytes} bytes<br>
                    <strong>KV Cache Size:</strong> ${totalBytes} / (1024³) ≈ ${sizeGB.toFixed(4)} GB
                `;
            }
            
            resultDiv.className = 'result show';
            resultDiv.innerHTML = `
                <div class="result-value">KV Cache Size: ${sizeGB.toFixed(4)} GB</div>
                <div class="result-label">Calculation Details:</div>
                <div class="details">${details}</div>
            `;
        }
        
        // Calculate maximum tokens from GPU RAM
        function calculateMaxTokens() {
            const model = document.getElementById('model-reverse').value;
            const gpuRamGB = parseFloat(document.getElementById('gpu-ram').value);
            const dtype = document.getElementById('dtype-reverse').value;
            const resultDiv = document.getElementById('result-reverse');
            
            // Validation
            if (isNaN(gpuRamGB) || gpuRamGB <= 0) {
                resultDiv.className = 'result show error';
                resultDiv.innerHTML = '<div class="result-label">Error</div><div>Please enter a valid GPU RAM size (greater than 0).</div>';
                return;
            }
            
            const config = modelConfigs[model];
            if (!config) {
                resultDiv.className = 'result show error';
                resultDiv.innerHTML = '<div class="result-label">Error</div><div>Model configuration not found.</div>';
                return;
            }
            
            // Get dtype size
            const dtypeSizes = {
                'float32': 4,
                'float16': 2,
                'bfloat16': 2,
                'int8': 1
            };
            const dtypeSize = dtypeSizes[dtype];
            
            // Convert GPU RAM to bytes
            const totalBytes = gpuRamGB * (1024 ** 3);
            
            // Calculate based on model architecture
            const isDeepSeek = model.includes('DeepSeek');
            const isQwen3 = model.toLowerCase().includes('qwen/qwen3-');
            
            let maxTokens;
            let details = '';
            
            if (isDeepSeek && config.kv_lora_rank && config.qk_rope_head_dim) {
                // DeepSeek V3/R1 architecture: tokens = total_bytes / (layers × (kv_lora_rank + qk_rope_head_dim) × dtype_size)
                const elementsPerToken = config.num_hidden_layers * (config.kv_lora_rank + config.qk_rope_head_dim);
                const bytesPerToken = elementsPerToken * dtypeSize;
                maxTokens = Math.floor(totalBytes / bytesPerToken);
                
                details = `
                    <strong>Selected Model:</strong> ${model}<br>
                    <strong>GPU RAM Size:</strong> ${gpuRamGB} GB<br>
                    <strong>Number of Hidden Layers:</strong> ${config.num_hidden_layers}<br>
                    <strong>KV LoRA Rank:</strong> ${config.kv_lora_rank}<br>
                    <strong>QK RoPE Head Dim:</strong> ${config.qk_rope_head_dim}<br>
                    <strong>Data Type Size:</strong> ${dtypeSize} bytes<br>
                    <strong>Elements per Token:</strong> ${config.num_hidden_layers} × (${config.kv_lora_rank} + ${config.qk_rope_head_dim}) = ${elementsPerToken}<br>
                    <strong>Bytes per Token:</strong> ${elementsPerToken} × ${dtypeSize} = ${bytesPerToken} bytes<br>
                    <strong>Maximum Tokens:</strong> ${totalBytes} / ${bytesPerToken} = ${maxTokens} tokens
                `;
            } else if (isQwen3 && config.head_dim) {
                // Qwen3: tokens = total_bytes / (2 × layers × kv_heads × head_dim × dtype_size)
                const elementsPerToken = 2 * config.num_hidden_layers * config.num_key_value_heads * config.head_dim;
                const bytesPerToken = elementsPerToken * dtypeSize;
                maxTokens = Math.floor(totalBytes / bytesPerToken);
                
                details = `
                    <strong>Selected Model:</strong> ${model}<br>
                    <strong>GPU RAM Size:</strong> ${gpuRamGB} GB<br>
                    <strong>Hidden Size:</strong> ${config.hidden_size}<br>
                    <strong>Number of Attention Heads:</strong> ${config.num_attention_heads}<br>
                    <strong>Number of Hidden Layers:</strong> ${config.num_hidden_layers}<br>
                    <strong>Number of Key-Value Heads:</strong> ${config.num_key_value_heads}<br>
                    <strong>Head Dimension:</strong> ${config.head_dim}<br>
                    <strong>Data Type Size:</strong> ${dtypeSize} bytes<br>
                    <strong>Elements per Token:</strong> 2 × ${config.num_hidden_layers} × ${config.num_key_value_heads} × ${config.head_dim} = ${elementsPerToken}<br>
                    <strong>Bytes per Token:</strong> ${elementsPerToken} × ${dtypeSize} = ${bytesPerToken} bytes<br>
                    <strong>Maximum Tokens:</strong> ${totalBytes} / ${bytesPerToken} = ${maxTokens} tokens
                `;
            } else {
                // Standard transformer: tokens = total_bytes / (2 × layers × kv_heads × head_size × dtype_size)
                const headSize = config.hidden_size / config.num_attention_heads;
                const elementsPerToken = 2 * config.num_hidden_layers * config.num_key_value_heads * headSize;
                const bytesPerToken = elementsPerToken * dtypeSize;
                maxTokens = Math.floor(totalBytes / bytesPerToken);
                
                details = `
                    <strong>Selected Model:</strong> ${model}<br>
                    <strong>GPU RAM Size:</strong> ${gpuRamGB} GB<br>
                    <strong>Hidden Size:</strong> ${config.hidden_size}<br>
                    <strong>Number of Attention Heads:</strong> ${config.num_attention_heads}<br>
                    <strong>Number of Hidden Layers:</strong> ${config.num_hidden_layers}<br>
                    <strong>Number of Key-Value Heads:</strong> ${config.num_key_value_heads}<br>
                    <strong>Head Size:</strong> ${headSize} (Hidden Size / Attention Heads)<br>
                    <strong>Data Type Size:</strong> ${dtypeSize} bytes<br>
                    <strong>Elements per Token:</strong> 2 × ${config.num_hidden_layers} × ${config.num_key_value_heads} × ${headSize} = ${elementsPerToken}<br>
                    <strong>Bytes per Token:</strong> ${elementsPerToken} × ${dtypeSize} = ${bytesPerToken} bytes<br>
                    <strong>Maximum Tokens:</strong> ${totalBytes} / ${bytesPerToken} = ${maxTokens} tokens
                `;
            }
            
            resultDiv.className = 'result show';
            resultDiv.innerHTML = `
                <div class="result-value">Maximum Tokens: ${maxTokens.toLocaleString()}</div>
                <div class="result-label">Calculation Details:</div>
                <div class="details">${details}</div>
            `;
        }
        
        // Enter key handler
        document.addEventListener('DOMContentLoaded', () => {
            const tokensInput = document.getElementById('tokens');
            if (tokensInput) {
                tokensInput.addEventListener('keypress', (e) => {
                    if (e.key === 'Enter') {
                        calculateKVCache();
                    }
                });
            }
            
            const gpuRamInput = document.getElementById('gpu-ram');
            if (gpuRamInput) {
                gpuRamInput.addEventListener('keypress', (e) => {
                    if (e.key === 'Enter') {
                        calculateMaxTokens();
                    }
                });
            }
        });
        
        // Load configurations on page load
        window.addEventListener('load', loadModelConfigs);
    </script>
</body>
</html>

